namespace NBody.Analysis
{
    using System;
    using System.Collections;
    using DataFile;
    
    public class Bullock
    {
        private struct Cell
        {
            public int N;
            public double Mass;
            public double Jx;
            public double Jy;
            public double Jz;
            public double J;
            public double Radius;
            public double EncMass;            
        }
        
        // sorts cells in reverse order, those with the most J at the top
        private class CellSorter : IComparer 
        {
            public int Compare( Object x, Object y )
            {
                double xx = ((Cell)x).J;
                double yy = ((Cell)y).J;
                if (xx < yy)
                    return 1;
                else if (yy < xx)
                    return -1;
                else
                    return 0;
            }
        }
            
        public static Table CalculateAngularMomentumDistribution(ref NBodySystem s, int num_per_bin)
        {
            Console.Error.WriteLine("Calculatig the angular momentum distribution");
            
            int num_bins = (int)(Math.Ceiling((double)s.NumParts / (double)num_per_bin));
            
            Cell[] cells = new Cell[num_bins * 3];
            Console.Error.WriteLine("Using {0} cells", cells.Length);
            
            s.Sort();
            
            // split into cells
            double vol1 = Math.Sqrt(5.0) / 3.0;
            double vol2 = Math.Sqrt(8.0) / 3.0;
            // system is sorted, so just put first n into bin 1, etc
            int c = 0;
            double enc_mass = 0.0;
            for (int i = 0; i < s.NumParts; i++)
            {
                double jz = s[i].Mass * s[i].Jz;
                double sin = Math.Sin(s[i].Theta);
                enc_mass += s[i].Mass;
                if (sin < vol1)
                {
                    cells[c + 0].N++;
                    cells[c + 0].J += jz;
                    cells[c + 0].Mass += s[i].Mass;
                    cells[c + 0].Radius = ((cells[c+0].N - 1.0) * cells[c+0].Radius + s[i].Radius) / cells[c+0].N;
                    cells[c + 0].EncMass = enc_mass; 
                }
                else if (sin < vol2)
                {
                    cells[c + 1].N++;
                    cells[c + 1].J += jz;
                    cells[c + 1].Mass += s[i].Mass;
                    cells[c + 1].Radius = ((cells[c+1].N - 1.0) * cells[c+1].Radius + s[i].Radius) / cells[c+1].N;
                    cells[c + 1].EncMass = enc_mass; 
                }
                else
                {
                    cells[c + 2].N++;
                    cells[c + 2].J += jz;
                    cells[c + 2].Mass += s[i].Mass;
                    cells[c + 2].Radius = ((cells[c+2].N - 1.0) * cells[c+2].Radius + s[i].Radius) / cells[c+2].N;
                    cells[c + 2].EncMass = enc_mass; 
                }
            
                if ((i+1) % num_per_bin == 0)
                    c += 3;
            }
            
            // how many cells have negative jz?
            int c_neg = 0;
            foreach (Cell x in cells)
            {
                if (x.J < 0)
                    c_neg++;
            }
            double frac_neg = (double)c_neg / (double)cells.Length;
            Console.Error.WriteLine("{0:f2}% of cells are negatively aligned", 100.0 * frac_neg);
            
            // hmm ... I think having "specific" angular momentum should come in here:
            for (int i = 0; i < cells.Length; i++)
            {
                if (cells[i].Mass > 0.0)
                    cells[i].J /= cells[i].Mass;
            }
            
            // sort cells by j, largest at the top.
            IComparer comp = new CellSorter();
            Array.Sort(cells, comp);
            
            DataFile.Table data = new DataFile.Table(cells.Length, 4, "Specific angular momentum (J, M(<J), Jmin, Jmax)");
            double jmax = cells[0].J;
            double mmax = 0.0;
            for (int i = 0; i < cells.Length; i++)
            {
                data[i, 0] = cells[i].J;
                for (int j = i + 1; j < cells.Length; j++)
                {
                    if (cells[j].J > 0)  // only include if positive
                    {
                        data[i, 1] += cells[j].Mass;
                    }
                }
                // error bars -- min and max J
                // not sure about this error formulation -- not exactly like in Bullock et al
                double error = 0.2 * cells[i].Radius * Math.Sqrt(cells[i].EncMass / cells[i].Radius) / Math.Sqrt(cells[i].N) / cells[i].Mass;
                //double error = 1.0 / Math.Sqrt(cells[i].N);
                //double error = 0.2 * Math.Sqrt(2.0 * cells[i].EncMass * cells[i].EncMass * cells[i].EncMass * cells[i].Radius / (double)cells[i].N) / cells[i].Mass;
                data[i, 2] = data[i, 0] - error;
                data[i, 3] = data[i, 0] + error;
                if (cells[i].J > jmax) jmax = cells[i].J;
                if (cells[i].J > 0.0) mmax += cells[i].Mass;
            }
            
            // now make the data relative
            Console.WriteLine("{0} {1} {2}", s.Time, jmax, mmax);
            for (int i = 0; i < data.Length; i++)
            {
                //Console.Error.WriteLine("{0} {1}", data[i, 0], data[i, 1]);
                data[i, 0] /= jmax;
                data[i, 1] /= mmax;
                data[i, 2] /= jmax;
                data[i, 3] /= jmax;
            }
            Console.Error.WriteLine("Jmax = {0}", jmax);
        
            return data;
        }
        
        public static Table CalculateAngularMomentumDistributionJ2(ref NBodySystem s, int num_per_bin)
        {
            Console.Error.WriteLine("Calculatig the angular momentum distribution for J^2");
            
            int num_bins = (int)(Math.Ceiling((double)s.NumParts / (double)num_per_bin));
            
            Cell[] cells = new Cell[num_bins * 3];
            Console.Error.WriteLine("Using {0} cells", cells.Length);
            
            s.Sort();
            
            // split into cells
            double vol1 = Math.Sqrt(5.0) / 3.0;
            double vol2 = Math.Sqrt(8.0) / 3.0;
            // system is sorted, so just put first n into bin 1, etc
            int c = 0;
            double enc_mass = 0.0;
            for (int i = 0; i < s.NumParts; i++)
            {
                double jx = s[i].Mass * s[i].Jx;
                double jy = s[i].Mass * s[i].Jy;
                double jz = s[i].Mass * s[i].Jz;
                double sin = Math.Sin(s[i].Theta);
                enc_mass += s[i].Mass;
                if (sin < vol1)
                {
                    cells[c + 0].N++;
                    cells[c + 0].Jx += jx;
                    cells[c + 0].Jy += jy;
                    cells[c + 0].Jz += jz;
                    cells[c + 0].Mass += s[i].Mass;
                    cells[c + 0].Radius = ((cells[c+0].N - 1.0) * cells[c+0].Radius + s[i].Radius) / cells[c+0].N;
                    cells[c + 0].EncMass = enc_mass; 
                }
                else if (sin < vol2)
                {
                    cells[c + 1].N++;
                    cells[c + 1].Jx += jx;
                    cells[c + 1].Jy += jy;
                    cells[c + 1].Jz += jz;
                    cells[c + 1].Mass += s[i].Mass;
                    cells[c + 1].Radius = ((cells[c+1].N - 1.0) * cells[c+1].Radius + s[i].Radius) / cells[c+1].N;
                    cells[c + 1].EncMass = enc_mass; 
                }
                else
                {
                    cells[c + 2].N++;
                    cells[c + 1].Jx += jx;
                    cells[c + 1].Jy += jy;
                    cells[c + 1].Jz += jz;
                    cells[c + 2].Mass += s[i].Mass;
                    cells[c + 2].Radius = ((cells[c+2].N - 1.0) * cells[c+2].Radius + s[i].Radius) / cells[c+2].N;
                    cells[c + 2].EncMass = enc_mass; 
                }
            
                if ((i+1) % num_per_bin == 0)
                    c += 3;
            }
            
            for (int i = 0; i < cells.Length; i++)
                cells[i].J = Math.Sqrt(cells[i].Jx*cells[i].Jx + cells[i].Jy*cells[i].Jy + cells[i].Jz*cells[i].Jz);
            
            // hmm ... I think having "specific" angular momentum should come in here:
            for (int i = 0; i < cells.Length; i++)
            {
                if (cells[i].Mass > 0.0)
                    cells[i].J /= cells[i].Mass;
            }
            
            // sort cells by j, largest at the top.
            IComparer comp = new CellSorter();
            Array.Sort(cells, comp);
            
            DataFile.Table data = new DataFile.Table(cells.Length, 4, "Specific angular momentum (J2, M(<J), Jmin, Jmax)");
            double jmax = cells[0].J;
            double mmax = 0.0;
            for (int i = 0; i < cells.Length; i++)
            {
                data[i, 0] = cells[i].J;
                for (int j = i + 1; j < cells.Length; j++)
                {
                    if (cells[j].J > 0)  // only include if positive
                    {
                        data[i, 1] += cells[j].Mass;
                    }
                }
                // error bars -- min and max J
                // not sure about this error formulation -- not exactly like in Bullock et al
                double error = 0.2 * cells[i].Radius * Math.Sqrt(cells[i].EncMass / cells[i].Radius) / Math.Sqrt(cells[i].N) / cells[i].Mass;
                //double error = 1.0 / Math.Sqrt(cells[i].N);
                //double error = 0.2 * Math.Sqrt(2.0 * cells[i].EncMass * cells[i].EncMass * cells[i].EncMass * cells[i].Radius / (double)cells[i].N) / cells[i].Mass;
                data[i, 2] = data[i, 0] - error;
                data[i, 3] = data[i, 0] + error;
                if (cells[i].J > jmax) jmax = cells[i].J;
                if (cells[i].J > 0.0) mmax += cells[i].Mass;
            }
            
            // now make the data relative
            for (int i = 0; i < data.Length; i++)
            {
                data[i, 0] /= jmax;
                data[i, 1] /= mmax;
                data[i, 2] /= jmax;
                data[i, 3] /= jmax;
            }
        
            return data;
        }
        
        public static void CalculateTotalJVector(NBodySystem s, out double jx, out double jy, out double jz)
        {
            jx = jy = jz = 0.0;
            foreach (Particle p in s)
            {
                jx += p.Mass * (p.Y * p.Vz - p.Z * p.Vy);
                jy += p.Mass * (p.Z * p.Vx - p.X * p.Vz);
                jz += p.Mass * (p.X * p.Vy - p.Y * p.Vx);
            }
        }
        
        public static double CalculateAlignment(ref NBodySystem s)
        {
            s.Sort();
            
            // find the radius that encloses half the mass
            double m_tot = s.Mass;
            double m_half = m_tot / 2.0;
        
            double m_check = 0.0;
            int half = 0;
            while (m_check < m_half)
                m_check += s[half++].Mass;
        
            // now calculate j within that radius
            double jx1, jy1, jz1;
            jx1 = jy1 = jz1 = 0.0;
            for (int i = 0; i < half; i++)
            {
                jx1 += s[i].Mass * (s[i].Y * s[i].Vz - s[i].Z * s[i].Vy);
                jy1 += s[i].Mass * (s[i].Z * s[i].Vx - s[i].X * s[i].Vz);
                jz1 += s[i].Mass * (s[i].X * s[i].Vy - s[i].Y * s[i].Vx);
            }
            double abs1 = Math.Sqrt(jx1 * jx1 + jy1 * jy1 + jz1 * jz1);
            jx1 /= abs1;
            jy1 /= abs1;
            jz1 /= abs1;
            Console.Error.WriteLine("Inner halo (r < {0:f4}):", s[half].Radius);
            Console.Error.WriteLine("  jx = {0:f4}\n  jy = {1:f4}\n  jz = {2:f4}", jx1, jy1, jz1);
        
            // now do the other half
            double jx2, jy2, jz2;
            jx2 = jy2 = jz2 = 0.0;
            for (int i = half; i < s.NumParts; i++)
            {
                jx2 += s[i].Mass * (s[i].Y * s[i].Vz - s[i].Z * s[i].Vy);
                jy2 += s[i].Mass * (s[i].Z * s[i].Vx - s[i].X * s[i].Vz);
                jz2 += s[i].Mass * (s[i].X * s[i].Vy - s[i].Y * s[i].Vx);
            }
            double abs2 = Math.Sqrt(jx2 * jx2 + jy2 * jy2 + jz2 * jz2);
            jx2 /= abs2;
            jy2 /= abs2;
            jz2 /= abs2;
            Console.Error.WriteLine("Out halo (r > {0:f4}):", s[half].Radius);
            Console.Error.WriteLine("  jx = {0:f4}\n  jy = {1:f4}\n  jz = {2:f4}", jx2, jy2, jz2);
        
            double cos = (jx1 * jx2 + jy1 * jy2 + jz1 * jz2);
        
            return cos;
        }
        
        public static DataFile.Table CalculateShells(NBodySystem s, int num_per_bin)
        {
            int num_bins = (int)(Math.Ceiling((double)s.NumParts / (double)num_per_bin));
            double rmax = s.Rmax + 1e-10;
            double rmin = s.Rmin;
            double deltaR = Math.Log10(rmax / rmin) / (double)num_bins;
        
            double[] mass = new double[num_bins];
            double[] jxbin = new double[num_bins];
            double[] jybin = new double[num_bins];
            double[] jzbin = new double[num_bins];
        
            DataFile.Table data = new DataFile.Table(num_bins, 5, "Angular momentum in spherical shells");
            for (int i = 0; i < num_bins; i++)
            {
                jxbin[i] = jybin[i] = jzbin[i] = 0.0;
                data[i, 0] = 0.0; // enclosed mass
                data[i, 1] = 0.0; // average angular momentum in shell
                data[i, 2] = 0.0;
                mass[i] = 0.0; // mass in each shell
            }
        
            for (int i = 0; i < s.NumParts; i++)
            {
                int pos = (int)((Math.Log10(s[i].Radius) - Math.Log10(rmin)) / deltaR);
                jxbin[pos] += s[i].Mass * s[i].Jx;
                jybin[pos] += s[i].Mass * s[i].Jy;
                jzbin[pos] += s[i].Mass * s[i].Jz;
                mass[pos] += s[i].Mass;
            }
        
            double m_enc = 0.0;
            double jmax = 0.0;
            for (int i = 0; i < num_bins; i++)
            {
                data[i, 0] = Math.Pow(10.0, Math.Log10(rmin) + 0.5 * deltaR + deltaR * i);
                m_enc += mass[i];
                data[i, 1] = m_enc;
                if (mass[i] > 0.0)
                    data[i, 2] = Math.Sqrt(jxbin[i] * jxbin[i] + jybin[i] * jybin[i] + jzbin[i] * jzbin[i]);
                else
                    data[i, 2] = 0.0;
                if (data[i, 2] > jmax) jmax = data[i, 2];
                
                data[i, 3] = Math.Sqrt(m_enc * data[i, 0]);
                data[i, 4] = Math.Sqrt(mass[i] * data[i, 0]);
            }
        
            // normalize data ?
            /*for (int i = 0; i < num_bins; i++)
            {
                data[i, 1] /= m_enc;
                data[i, 2] /= jmax;
            }*/
        
            return data;
        }
        
        public static DataFile.Table CalculateShells2(NBodySystem s, int num_per_bin)
        {
            s.Sort();
            
            int n = s.NumParts / (num_per_bin);
        
            DataFile.Table data = new DataFile.Table(n, 2, "Angular momentum in spherical shells");
            double jx = 0.0, jy = 0.0, jz = 0.0;
            double mass = 0.0;
            double mbin = 0.0;
            int c = 0;
            for (int i = 0; i < s.NumParts; i++)
            {
                jx += s[i].Jx;
                jy += s[i].Jy;
                jz += s[i].Jz;
                mass += s[i].Mass;
                mbin += s[i].Mass;
                if ((i != 0 && i % num_per_bin == 0) || i == s.NumParts)
                {
                    data[c, 0] = mass;
                    data[c, 1] = Math.Sqrt(jx*jx + jy*jy + jz*jz);
                    jx = jy = jz = 0.0;
                    mbin = 0.0;
                    c++;
                }
            }
            
            double jmax = 0.0;
            for (int i = 0; i < n; i++)
                if (data[i, 1] > jmax) jmax = data[i, 1];
            
            for (int i = 0; i < n; i++)
            {
                data[i, 0] /= mass;
                data[i, 1] /= jmax;
            }
            
            return data;            
        }
    }
}
